package com.abc.doc;

import cn.hutool.core.util.ObjectUtil;
import com.abc.dto.response.RerankResponseDto;
import com.abc.service.VectorStoreService;
import com.abc.util.RRFUtil;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

@Slf4j
@Component
public class KnowledgeDocumentPostProcessor implements DocumentPostProcessor {

    @Autowired
    private VectorStoreService vectorStoreService;

    @Override
    public List<Document> process(Query query, List<Document> documents) {
        RRFUtil.enhanceRanking(documents, query.text());
        RerankResponseDto rerankResponseDto = vectorStoreService.performRerank(query.text(), documents);
        if (ObjectUtil.isNotNull(rerankResponseDto)) {
            return rerankResponseDto.getResults().stream().filter(ObjectUtil::isNotEmpty).map(r -> new Document(r.getDocument().getText())).toList();
        } else {
            return documents;
        }
    }
}
